package com.atguigu.day06;

import com.atguigu.utils.ClickEvent;
import com.atguigu.utils.ClickSource;
import com.atguigu.utils.UrlViewCountPerWindow;
import org.apache.flink.api.common.eventtime.*;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

// 自定义水位线发生器
public class Example1 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment
                .getExecutionEnvironment();
        env.setParallelism(1);

        env
                .addSource(new ClickSource())
                .assignTimestampsAndWatermarks(new CustomWatermark())
                .keyBy(r -> r.url)
                .window(TumblingEventTimeWindows.of(Time.seconds(5)))
                .aggregate(new CountAgg(), new WindowResult())
                .print();

        env.execute();
    }

    public static class CustomWatermark implements WatermarkStrategy<ClickEvent> {
        @Override
        public TimestampAssigner<ClickEvent> createTimestampAssigner(TimestampAssignerSupplier.Context context) {
            return new SerializableTimestampAssigner<ClickEvent>() {
                @Override
                public long extractTimestamp(ClickEvent element, long recordTimestamp) {
                    return element.ts; // 指定事件时间是哪一个字段
                }
            };
        }

        @Override
        public WatermarkGenerator<ClickEvent> createWatermarkGenerator(WatermarkGeneratorSupplier.Context context) {
            return new WatermarkGenerator<ClickEvent>() {
                long bound = 5000L; // 最大延迟时间
                // 为了防止溢出，加bound和1毫秒
                long maxTs = -Long.MAX_VALUE + bound + 1L;
                // `onEvent`方法：每来一个事件调用一次
                @Override
                public void onEvent(ClickEvent event, long eventTimestamp, WatermarkOutput output) {
                    // 更新观察到的最大事件时间
                    maxTs = Math.max(event.ts, maxTs);
                    // 还可以在onEvent里发送水位线
                    // 这个水位线插入在event事件的后面
                    // output.emitWatermark(new Watermark(event.ts - 1L));
                }

                // `onPeriodicEmit`方法：周期性调用，默认每隔200毫秒调用一次
                @Override
                public void onPeriodicEmit(WatermarkOutput output) {
                    // 插入水位线
                    output.emitWatermark(new Watermark(
                            // 水位线 = 观察到的最大事件时间 - 最大延迟时间 - 1毫秒
                            maxTs - bound - 1L
                    ));
                }
            };
        }
    }

    public static class CountAgg implements AggregateFunction<ClickEvent, Long, Long> {
        @Override
        public Long createAccumulator() {
            return 0L;
        }

        @Override
        public Long add(ClickEvent value, Long accumulator) {
            return accumulator + 1L;
        }

        @Override
        public Long getResult(Long accumulator) {
            return accumulator;
        }

        @Override
        public Long merge(Long a, Long b) {
            return null;
        }
    }

    public static class WindowResult extends ProcessWindowFunction<Long, UrlViewCountPerWindow, String, TimeWindow> {
        @Override
        public void process(String s, Context context, Iterable<Long> elements, Collector<UrlViewCountPerWindow> out) throws Exception {
            out.collect(new UrlViewCountPerWindow(
                    s,
                    elements.iterator().next(),
                    context.window().getStart(),
                    context.window().getEnd()
            ));
        }
    }
}
